import jieba
from wordcloud import WordCloud
import traceback
import os
import platform


# 读取弹幕文件
def read_danmu_file(file_name):
    with open(file_name, mode='r', encoding='utf8') as f:
        danmu = f.read()
        return danmu


# jieba分词
def jieba_cut(danmu):
    jieba.suggest_freq('有人吗', tune=True)
    jieba.suggest_freq('另一个', tune=True)
    jieba.suggest_freq('为什么', tune=True)
    jieba.suggest_freq('钉钉', tune=True)
    jieba.suggest_freq('中门对狙', tune=True)
    jieba.suggest_freq('阿巴瑟', tune=True)
    cut_list = jieba.lcut(danmu)
    return cut_list


# 生成词云图
def gen_wordcloud(cut_list, file_name):
    word_str = ' '.join(cut_list)
    if platform.system() == 'Windows':
        wc_settings = {
            'font_path': 'msyh.ttc',
            'width': 1024,
            'height': 768,
            'background_color': 'white',
            'max_words': 50,
        }
    else:
        wc_settings = {
            'font_path': '/System/Library/fonts/PingFang.ttc',
            'width': 1024,
            'height': 768,
            'background_color': 'white',
            'max_words': 50,
        }
    wc = WordCloud(**wc_settings).generate(word_str)

    # 保存到图片
    tmp_name = file_name.rstrip('.txt')
    wc.to_file(f'{tmp_name}.png')


def file_list(file_dir):
    for root, dirs, files in os.walk(file_dir):
        print(f'路径：{root}')
        print(f'目录：{dirs}')
        print(f'文件：{files}')


def main(av_or_bv_number):
    try:
        file_name = f'{av_or_bv_number}.txt'
        danmu = read_danmu_file(file_name)
        cut_list = jieba_cut(danmu)
        gen_wordcloud(cut_list, file_name)
    except Exception:
        print(traceback.format_exc())


if __name__ == '__main__':
    pwd = os.getcwd()
    file_list(pwd)
    av = input('请输入要生成词云的弹幕文件名：')
    main(av)
